"""
This is function description
"""
import random

import numpy as np
import torch
import matplotlib.pyplot as plt

from environment.electric_scheduling import PowerDayAheadSchedule
from model.DDPG.ddpg_learning import DDPG
import utils.rl_utils as rl
import model.DDPG.ddpg_train_off_policy as top

actor_lr = 3e-4
critic_lr = 3e-3
num_episodes = 2000
hidden_dim = 128
gamma = 0.98
tau = 0.005  # 软更新参数
buffer_size = 10000
minimal_size = 1000
batch_size = 64
sigma = 0.01  # 高斯噪声标准差
device = "cuda" if torch.cuda.is_available() else "cpu"


def Trainer(env: PowerDayAheadSchedule):
    state, _ = env.reset()
    print(state)
    state_size = len(state)
    action_size = env.fire_station_num + env.water_station_num + 2 + 2
    agent = DDPG(
        state_dim=state_size,
        hidden_dim=hidden_dim,
        action_dim=action_size,
        action_bound=400,
        sigma=sigma,
        actor_lr=actor_lr,
        critic_lr=critic_lr,
        tau=tau,
        gamma=gamma,
        device=device
    )
    buffer = rl.ReplayBuffer(buffer_size)
    random.seed(0)
    np.random.seed(0)
    torch.manual_seed(0)
    return_list = top.train_off_policy_agent(env, agent, num_episodes, buffer, minimal_size, batch_size)
    episode_list = list(range(len(return_list)))
    plt.plot(episode_list, return_list)
    plt.xlabel('Episodes')
    plt.ylabel('Returns')
    plt.show()
